/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 集合通信算子build graph，用于构造执行图中的const args kernel和launch kernel
 * Author: lilianlin
 * Create: 2022-07-15
 */

#include "hcom_build_graph.h"
#include "hcom_op_utils.h"
#include "comm.h"
#include <map>
#include "adapter_qos_pub.h"
#include "graph/debug/ge_attr_define.h"
#include "exe_graph/lowering/dev_mem_value_holder.h"
#include <unordered_map>

using namespace std;

namespace hccl {
const std::string HCOM_ATTR_ADDR_LENGTH = "addr_length";
std::map<const std::string, HcomOpType> HcomOpTypeTransform = {
    { HCCL_KERNEL_OP_TYPE_ALLGATHER, HcomOpType::HCOM_ALL_GATHER },
    { HCCL_KERNEL_OP_TYPE_ALLGATHERV, HcomOpType::HCOM_ALL_GATHER_V },
    { HCCL_KERNEL_OP_TYPE_ALLREDUCE, HcomOpType::HCOM_ALL_REDUCE },
    { HCCL_KERNEL_OP_TYPE_BROADCAST, HcomOpType::HCOM_BROADCAST },
    { HCCL_KERNEL_OP_TYPE_REDUCESCATTER, HcomOpType::HCOM_REDUCE_SCATTER },
    { HCCL_KERNEL_OP_TYPE_REDUCESCATTERV, HcomOpType::HCOM_REDUCE_SCATTER_V },
    { HCCL_KERNEL_OP_TYPE_ALLTOALLV, HcomOpType::HCOM_ALL_TO_ALL_V },
    { HCCL_KERNEL_OP_TYPE_ALLTOALLVC, HcomOpType::HCOM_ALL_TO_ALL_VC },
    { HCCL_KERNEL_OP_TYPE_GATHER_ALLTOALLV, HcomOpType::HCOM_GATHER_ALL_TO_ALL_V },
    { HCCL_KERNEL_OP_TYPE_REMOTE_READ, HcomOpType::HCOM_REMOTE_READ },
    { HCCL_KERNEL_OP_TYPE_REMOTE_REF_READ, HcomOpType::HCOM_REMOTE_REF_READ },
    { HCCL_KERNEL_OP_TYPE_REMOTE_WRITE, HcomOpType::HCOM_REMOTE_WRITE },
    { HCCL_KERNEL_OP_TYPE_REMOTE_SCATTER_WRITE, HcomOpType::HCOM_REMOTE_SCATTER_WRITE },
    { HCCL_KERNEL_OP_TYPE_SEND, HcomOpType::HCOM_SEND },
    { HCCL_KERNEL_OP_TYPE_RECEIVE, HcomOpType::HCOM_RECEIVE },
    { HCCL_KERNEL_OP_TYPE_REDUCE, HcomOpType::HCOM_REDUCE },
    { HCCL_KERNEL_OP_TYPE_REMOTE_LOOKUP, HcomOpType::REMOTE_LOOK_UP },
    { HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_PAIRED, HcomOpType::COLL_REMOTE_LOOK_UP_PAIRED },
    { HCCL_KERNEL_OP_TYPE_COLL_REMOTE_UPDATE_PAIRED, HcomOpType::COLL_REMOTE_UPDATE_PAIRED },
    { HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_UNIQUED_PAIRED, HcomOpType::COLL_REMOTE_LOOK_UP_UNIQUED_PAIRED },
};

HcclResult GetOpIntAttr(const ge::OpDescPtr &op, const string &attr, s32 &output)
{
        if (ge::AttrUtils::HasAttr(op, attr)) {
            ge::AttrUtils::GetInt(op, attr, output);
            HCCL_INFO("[Get][%s]:success. %s value[%d]", attr.c_str(), attr.c_str(), output);
        } else {
            HCCL_ERROR("[Get][%s]:failed. opDesc has no %s", attr.c_str(), attr.c_str());
            return HCCL_E_NOT_FOUND;
        }
        return HCCL_SUCCESS;
}

/*
 * **********************************************************************
 * 读取node中的属性，输出到属性的结构体中
 * 返回各算子的非公共属性
 * **********************************************************************
 */
HcclResult HcomAllGatherGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HcomOpUtils::GetRankSize(node->GetOpDesc(), opAttr.op.allgather.rankSize);
}

HcclResult HcomAllGatherVGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult HcomAllReduceGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HcomOpUtils::GetReduction(node->GetOpDesc(), opAttr.op.allreduce.reduction);
}

HcclResult HcomBroadcastGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HcomOpUtils::GetRoot(node->GetOpDesc(), opAttr.op.broadcast.root);
}

HcclResult HcomReduceScatterGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(HcomOpUtils::GetRankSize(node->GetOpDesc(), opAttr.op.reducescatter.rankSize));
    return HcomOpUtils::GetReduction(node->GetOpDesc(), opAttr.op.reducescatter.reduction);
}

HcclResult HcomReduceScatterVGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(HcomOpUtils::GetReduction(node->GetOpDesc(), opAttr.op.reducescatterv.reduction));
    return HCCL_SUCCESS;
}

HcclResult HcomAllToAllVGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    ge::GeTensorDesc outputTensor = node->GetOpDesc()->GetOutputDesc(0);
    ge::DataType geDataType = outputTensor.GetDataType();
    CHK_RET(HcomOpUtils::TransformDataType(geDataType, opAttr.op.alltoallv.recvType));
    return HCCL_SUCCESS;
}

HcclResult HcomAllToAllVCGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult HcomGatherAllToAllVGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_PRT_RET(
        (!ge::AttrUtils::GetInt(node->GetOpDesc(), HCOM_ATTR_ADDR_LENGTH, opAttr.op.gatheralltoallv.addrLength)),
        HCCL_ERROR("[Get][addrLength] get attr[%s] failed.", HCOM_ATTR_ADDR_LENGTH.c_str()), HCCL_E_PARA);
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteReadGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteRefReadGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult HcomRemoteWriteGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult RemoteScatterWriteGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    return HCCL_SUCCESS;
}

HcclResult HcomRecvGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(HcomOpUtils::GetSrcRank(node->GetOpDesc(), opAttr.op.recv.srcRank));
    return HcomOpUtils::GetSrTag(node->GetOpDesc(), opAttr.op.recv.srTag);
}

HcclResult HcomSendGetOpAttr(const ge::NodePtr &node, struct HcomOpAttr &opAttr)
{
    CHK_RET(HcomOpUtils::GetDestRank(node->GetOpDesc(), opAttr.op.send.destRank));
    return HcomOpUtils::GetSrTag(node->GetOpDesc(), opAttr.op.send.srTag);
}

HcclResult HcomReduceGetOpAttr(const ge::NodePtr &node, struct HcomOpAttr &opAttr)
{
    CHK_RET(HcomOpUtils::GetReduction(node->GetOpDesc(), opAttr.op.reduce.reduction));
    CHK_RET(HcomOpUtils::GetRoot(node->GetOpDesc(), opAttr.op.reduce.root));
    return HCCL_SUCCESS;
}

HcclResult RemoteLookupGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "tag", opAttr.op.remotelookup.tag));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "insert_option", opAttr.op.remotelookup.insert_option));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "max_num", opAttr.op.remotelookup.max_num));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "embedding_dim", opAttr.op.remotelookup.embedding_dim));
    return HCCL_SUCCESS;
}
 
HcclResult CollRemoteLookupPairedGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "tag", opAttr.op.collremotelookuppaired.tag));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "insert_option", opAttr.op.collremotelookuppaired.insert_option));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "max_num", opAttr.op.collremotelookuppaired.max_num));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "embedding_dim", opAttr.op.collremotelookuppaired.embedding_dim));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "flags", opAttr.op.collremotelookuppaired.flags));
    return HCCL_SUCCESS;
}
 
HcclResult CollRemoteUpdatePairedGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "tag", opAttr.op.collremoteupdatepaired.tag));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "max_num", opAttr.op.collremoteupdatepaired.max_num));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "embedding_dim", opAttr.op.collremoteupdatepaired.embedding_dim));
    return HCCL_SUCCESS;
}

HcclResult CollRemoteLookupUniquedAndPairedGetOpAttr(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "tag", opAttr.op.collremotelookuppairedandunique.tag));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "insert_option", opAttr.op.collremotelookuppairedandunique.insert_option));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "max_num", opAttr.op.collremotelookuppairedandunique.max_num));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "embedding_dim", opAttr.op.collremotelookuppairedandunique.embedding_dim));
    CHK_RET(GetOpIntAttr(node->GetOpDesc(), "flags", opAttr.op.collremotelookuppairedandunique.flags));
    return HCCL_SUCCESS;
}

std::vector<std::function<HcclResult(const ge::NodePtr&, struct HcomOpAttr&)>> HcomOpGetAttrFuncs = {
    HcomAllGatherGetOpAttr,
    HcomAllGatherVGetOpAttr,
    HcomAllReduceGetOpAttr,
    HcomBroadcastGetOpAttr,
    HcomReduceScatterGetOpAttr,
    HcomReduceScatterVGetOpAttr,
    HcomAllToAllVGetOpAttr,
    HcomAllToAllVCGetOpAttr,
    HcomGatherAllToAllVGetOpAttr,
    HcomRemoteReadGetOpAttr,
    HcomRemoteRefReadGetOpAttr,
    HcomRemoteWriteGetOpAttr,
    RemoteScatterWriteGetOpAttr,
    HcomSendGetOpAttr,
    HcomReduceGetOpAttr,
    CollRemoteLookupPairedGetOpAttr,
    CollRemoteUpdatePairedGetOpAttr,
    CollRemoteLookupUniquedAndPairedGetOpAttr,
    RemoteLookupGetOpAttr,
    HcomRecvGetOpAttr,
};

/*
 * **********************************************************************
 * 读取node中的属性，输出到属性的结构体中
 * 属性包含：opType，dataType，group，算子的非公共属性
 * **********************************************************************
 */
HcclResult GenerateHcomOpArg(const ge::NodePtr& node, struct HcomOpAttr& opAttr)
{
    auto iter = HcomOpTypeTransform.find(node->GetOpDesc()->GetType());
    if (iter == HcomOpTypeTransform.end()) {
        HCCL_ERROR("[Generate][HcomOpAttr] %s is not supported.", node->GetOpDesc()->GetType().c_str());
        return HCCL_E_PARA;
    }
    opAttr.opType = iter->second;
    auto opDesc = node->GetOpDesc();
    CHK_RET(HcomOpUtils::GetDataType(opDesc, opAttr.dataType));

    u32 groupId = HCCL_STREAM_DEFAULT_GROUP_ID;
    if (ge::AttrUtils::HasAttr(opDesc, ge::ATTR_NAME_PARALLEL_GROUP_ID)) {
        if (ge::AttrUtils::GetInt(opDesc, ge::ATTR_NAME_PARALLEL_GROUP_ID, groupId) == false) {
            HCCL_ERROR("[Get][QosConfig]errNo[0x%016llx]: get qosCfg failed. get \"qosCfg\" from"
                "opDesc failed", HCOM_ERROR_CODE(HCCL_E_PARA));
            return HCCL_E_PARA;
        }
        HCCL_DEBUG("Call GenerateHcomOpArg, groupId value[%x]", groupId);
    }

    u32 qosCfg = 0xFFFFFFFFU; // qos不使能的情况下为全F
    CHK_RET(hrtGetQosConfig(groupId, qosCfg));
    opAttr.qosConfig = qosCfg;

    std::string sGroup;
    CHK_RET(HcomOpUtils::GetGroup(opDesc, sGroup));
    int32_t sret = memcpy_s(&opAttr.group[0], sizeof(opAttr.group), sGroup.c_str(), (sGroup.length() + 1));
    CHK_PRT_RET(sret != EOK, HCCL_ERROR("[Generate][HcomOpAttr]memcpy failed. ret[%d],"
        "params:destMaxSize[%zu],count[%zu]", sret, sizeof(opAttr.group),
        (sGroup.length() + 1)), HCCL_E_SYSCALL);

    CHK_RET(HcomOpGetAttrFuncs[static_cast<int32_t>(opAttr.opType)](node, opAttr));
    return HCCL_SUCCESS;
}

/*
 * **********************************************************************
 * 执行图构图，集合通信算子node在lowering过程中，const args节点的构造
 * 包含算子的const options
 * **********************************************************************
 */
gert::bg::ValueHolderPtr GenerateHcomOpArgs(const ge::NodePtr& node)
{
    struct HcomOpAttr opAttr;
    HcclResult ret = GenerateHcomOpArg(node, opAttr);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("Node[%s]: Generate ConstOpArgs failed.", node->GetName().c_str());
        return nullptr;
    }
    return gert::bg::ValueHolder::CreateConst(&opAttr, sizeof(struct HcomOpAttr), false);
}

/*
 * **********************************************************************
 * 执行图构图，集合通信算子node在lowering过程中，launch kernel节点的构造
 * 包含算子的inputs，outputs，options
 * **********************************************************************
 */
std::vector<gert::bg::ValueHolderPtr> LaunchHcomOpKernel(const HcomLaunchArg &args,
    const std::vector<gert::bg::DevMemValueHolderPtr> &inputAddrs,
    const std::vector<gert::bg::DevMemValueHolderPtr> &outputAddrs,
    const std::vector<gert::bg::ValueHolderPtr> &inputShapes,
    const std::vector<gert::bg::ValueHolderPtr> &outputShapes)
{
    std::vector<gert::bg::ValueHolderPtr> inputs;
    uint32_t inputNum = inputAddrs.size();
    uint32_t outputNum = outputAddrs.size();
    inputs.emplace_back(args.opArgs);
    inputs.emplace_back(args.stream);
    inputs.emplace_back(gert::bg::ValueHolder::CreateConst(&inputNum, sizeof(inputNum)));
    inputs.emplace_back(gert::bg::ValueHolder::CreateConst(&outputNum, sizeof(outputNum)));
    inputs.insert(inputs.end(), inputAddrs.begin(), inputAddrs.end());
    inputs.insert(inputs.end(), outputAddrs.begin(), outputAddrs.end());
    inputs.insert(inputs.end(), inputShapes.begin(), inputShapes.end());
    inputs.insert(inputs.end(), outputShapes.begin(), outputShapes.end());
    return gert::bg::ValueHolder::CreateDataOutput("LaunchHcomKernel", inputs, outputAddrs.size());
}

std::vector<gert::bg::DevMemValueHolderPtr> LaunchRecvOpKernel(const HcomLaunchArg &args, const ge::NodePtr &node,
    const gert::LowerInput &lowerInput)
{
    std::vector<gert::bg::ValueHolderPtr> inputs;

    ge::DataType geDataType = node->GetOpDesc()->GetOutputDesc(0U).GetDataType();

    inputs.emplace_back(args.opArgs);
    inputs.emplace_back(args.stream);
    inputs.emplace_back(lowerInput.global_data->GetOrCreateAllocator({gert::kOnDeviceHbm,
                                                gert::AllocatorUsage::kAllocNodeOutput}));
    inputs.emplace_back(gert::bg::ValueHolder::CreateConst(&geDataType, sizeof(geDataType)));
    return gert::bg::DevMemValueHolder::CreateDataOutput("LaunchRecvKernel", inputs, LAUNCH_RECV_KERNEL_OUTPUT_SIZE,
        node->GetOpDesc()->GetStreamId());
}

}